// package: learn-go
// filename: wsChat
// author: diogoxiang@gmail.com
// time: 2024/1/9 14:14
package main

import (
	"github.com/gorilla/websocket"
	"log"
	"net/http"
)

type Message struct {
	Type    string
	Content string
	From    string
	To      string
}

func (m *Message) String() string {
	return m.Type + " " + m.Content + " From: " + m.From + " To: " + m.To
}

func ID() string {
	// 生成唯一ID的逻辑
	return "unique_id" // 示例中直接返回一个固定的唯一ID
}

type Client struct {
	conn *websocket.Conn
	ID   string
}

func (c *Client) Send(message Message) error {
	return c.conn.WriteMessage(websocket.TextMessage, []byte(message.String()))
}

func (c *Client) Receive() (Message, error) {
	var message Message
	err := websocket.Message.Receive(c.conn, &message)
	if err != nil {
		return Message{}, err
	}
	return message, nil
}

type Server struct {
	clients map[string]Client
}

func (s *Server) Register(client Client) {
	s.clients[client.ID] = client
}

func (s *Server) Unregister(client Client) {
	delete(s.clients, client.ID)
}

func (s *Server) Broadcast(message Message) {
	for _, client := range s.clients {
		if message.To == "" || client.ID == message.To {
			err := client.Send(message)
			if err != nil {
				log.Printf("Failed to send message to client %s: %v", client.ID, err)
			}
		}
	}
}

func main() {
	// 创建服务器
	server := Server{clients: make(map[string]Client)}

	// 监听 8080 端口
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// 这里应该处理静态文件，例如HTML
	})

	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 0, 0)
		if err != nil {
			log.Println(err)
			return
		}

		// 注册客户端
		client := Client{conn: conn, ID: ID()}
		server.Register(client)

		// 处理客户端消息
		go func() {
			for {
				message, err := client.Receive()
				if err != nil {
					log.Printf("Failed to receive message from client %s: %v", client.ID, err)
					server.Unregister(client)
					break
				}

				// 广播消息
				server.Broadcast(message)
			}
		}()
	})

	// 启动服务器
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatal(err)
	}
}
